import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import argparse
import os
from math import pi

def parse_args():
    parser = argparse.ArgumentParser(description="Generate Radar Chart for method comparison.")
    parser.add_argument('--files','-f', nargs='+', required=True, help='List of files (CSV or Excel) to analyze')
    parser.add_argument('--exclude', nargs='+', default=['BART'], help='Methods to exclude (default: BART)')
    parser.add_argument('--output','-o', default='radar_chart.png', help='Output filename for the plot')
    return parser.parse_args()

def read_data(file_path):
    _, ext = os.path.splitext(file_path)
    try:
        if ext.lower() in ['.xlsx', '.xls']:
            return pd.read_excel(file_path)
        else:
            return pd.read_csv(file_path)
    except Exception as e:
        print(f"Error reading {file_path}: {e}")
        return None

def get_aggregated_metrics(files, exclude_methods):
    """
    读取所有文件，计算每个方法在四个指标上的全局平均值。
    在此处进行更名操作。
    """
    all_data = []
    
    for f in files:
        df = read_data(f)
        if df is not None and not df.empty:
            # 过滤不需要的方法
            df = df[~df['method'].isin(exclude_methods)].copy()
            
            # === 核心修改：在这里进行重命名 ===
            df['method'] = df['method'].replace('ours(MTM)', 'ours(deepsignal3)')
            
            # 只保留需要的列
            cols = ['method', 'pearson', 'rsquare', 'spearman', 'RMSE']
            df = df[[c for c in cols if c in df.columns]]
            all_data.append(df)
            
    if not all_data:
        return None
        
    # 合并所有数据
    combined_df = pd.concat(all_data, ignore_index=True)
    
    # 按方法分组计算平均值
    agg_df = combined_df.groupby('method').mean().reset_index()
    return agg_df

def normalize_data(df):
    """
    将数据归一化到 [0, 1] 区间以便绘图。
    RMSE: 值越小越好 -> (Max - x) / Range -> 靠近1.0
    其他: 值越大越好 -> (x - Min) / Range -> 靠近1.0
    """
    df_norm = df.copy()
    metrics = ['pearson', 'rsquare', 'spearman', 'RMSE']
    
    for col in metrics:
        min_val = df[col].min()
        max_val = df[col].max()
        val_range = max_val - min_val
        
        if val_range == 0:
            df_norm[col] = 1.0 
        else:
            if col == 'RMSE':
                # RMSE: 越小越好，所以反转
                df_norm[col] = (max_val - df[col]) / val_range
            else:
                # 其他: 越大越好
                df_norm[col] = (df[col] - min_val) / val_range
                
    # 缩放区间到 [0.2, 1.0] 防止最小值缩在圆心
    for col in metrics:
        df_norm[col] = df_norm[col] * 0.8 + 0.2
        
    return df_norm

def plot_radar(agg_df, output_file):
    # 准备数据
    norm_df = normalize_data(agg_df)
    
    # 变量（指标）列表
    categories = ['Pearson', 'R-Square', 'Spearman', 'RMSE\n']
    N = len(categories)
    
    # 计算角度
    angles = [n / float(N) * 2 * pi for n in range(N)]
    angles += angles[:1] # 闭合回路
    
    # 初始化绘图
    fig, ax = plt.subplots(figsize=(8, 8), subplot_kw=dict(polar=True))
    
    # 设置方向：顺时针，且从12点钟方向开始
    ax.set_theta_offset(pi / 2)
    ax.set_theta_direction(-1)
    
    # 设置X轴标签（指标名称）
    plt.xticks(angles[:-1], categories, size=12, weight='bold')
    
    # 设置Y轴标签（刻度），隐藏具体数值，保持整洁
    ax.set_rlabel_position(0)
    plt.yticks([0.4, 0.6, 0.8, 1.0], ["", "", "", ""], color="grey", size=7)
    plt.ylim(0, 1.1)
    
    # === 定义样式：更新了名字 ===
    styles = {
        'Dorado': {
            'color': '#1f77b4', 'ls': '--', 'lw': 2, 'alpha': 0.1, 
            'label': 'Dorado (Gold Standard)'
        }, 
        'ours(deepsignal3)': {  # <--- 这里改成了新的名字
            'color': '#d62728', 'ls': '-', 'lw': 3, 'alpha': 0.2, 
            'label': 'Ours (deepsignal3)'
        },
        'Rockfish': {
            'color': '#2ca02c', 'ls': '-.', 'lw': 2, 'alpha': 0.05, 
            'label': 'Rockfish'
        },
        'DeepMod2': {
            'color': '#ff7f0e', 'ls': ':', 'lw': 2, 'alpha': 0.05, 
            'label': 'DeepMod2'
        }
    }
    
    # 绘图循环
    for index, row in norm_df.iterrows():
        method_name = row['method']
        values = row[['pearson', 'rsquare', 'spearman', 'RMSE']].values.flatten().tolist()
        values += values[:1] # 闭合
        
        # 获取样式
        style = styles.get(method_name, {'color': 'grey', 'ls': '-', 'lw': 1, 'alpha': 0.1, 'label': method_name})
        
        ax.plot(angles, values, linewidth=style['lw'], linestyle=style['ls'], label=style['label'], color=style['color'])
        ax.fill(angles, values, color=style['color'], alpha=style['alpha'])

    # 添加图例
    plt.legend(loc='upper right', bbox_to_anchor=(0.1, 0.1), fontsize=10)
    
    # 添加标题
    plt.title('Comprehensive Performance Comparison\n(Larger Area = Better Performance)', size=15, weight='bold', y=1.05)
    
    print(f"Saving radar chart to {output_file}")
    plt.savefig(output_file, dpi=300, bbox_inches='tight')
    plt.show()

def main():
    args = parse_args()
    
    print("Reading and aggregating data...")
    agg_df = get_aggregated_metrics(args.files, args.exclude)
    
    if agg_df is not None:
        print("Aggregated Data (Raw Means):")
        print(agg_df)
        
        plot_radar(agg_df, args.output)
    else:
        print("No data found to plot.")

if __name__ == "__main__":
    main()